from typing import Annotated, Union
from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect

from api.dependencies import get_cookie_or_token

router = APIRouter(
    prefix="/ws",
    tags=["ws"],
    # dependencies=[Depends(get_token_header)],
    responses={404: {"description": "Not found"}},
)


class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)


manager = ConnectionManager()


@router.websocket("/notice")
async def websocket_endpoint(
    *,
    websocket: WebSocket,
    client_id: int,
    q: Union[int, None] = None,
    cookie_or_token: Annotated[str, Depends(get_cookie_or_token)],
):
    # 记录用户与链接关系
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            await manager.send_personal_message(f"You wrote: {data}", websocket)
            # await manager.broadcast(f"Client #{client_id} says: {data}")
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        # 移除映射关系
        # await manager.broadcast(f"Client #{client_id} left the chat")
